# -*- coding: utf-8 -*-
"""
Created on Wed Aug 15 17:12:07 2012

@author: eduardo
"""

import numpy as np
import operator

#from scipy import sparse as spr


# Implement class for models
# Write documentation



def is_iterable(x):
    # Verifies if iterable
    try:
        iter(x)
    except TypeError:
        return False
    return True

class ADEMError(Exception):
    ''' Base class for error related to ADEM models'''
    pass
    

class DEMError(Exception):
    ''' Base class for errors related to DEM models. '''
    pass

class ModelError(DEMError):
    ''' Error produced during the specification of DEM model. '''
    def __init__(self,attribute,msg):
        ''' Attribute defines the error where the attribute was found.'''
        self.attribute = attribute
        self.msg = msg

class spm_vec_Error(DEMError):
    ''' Error while vectorizing.'''
    def __init__(self,msg=''):
        self.msg = msg
        
class spm_vec(object):
    '''Efficient vectorization of an arbitrary complex list of lists. 
    
    It is implemented as a container of inmutable type  and offers a large 
    number of operations on numeric arrays. It doesnt work on list of numbers
    as numbers are immutable.'''
    
    def __init__(self,iterable):

        if isinstance(iterable,spm_vec):
            return iterable
        
        def recursive_lecture(ai):
            a = np.array([])
            if isinstance(ai,np.ndarray):
                return ai.flat
            elif isinstance(ai,(list,tuple)):
                for i in ai:
                    r = recursive_lecture(i)
                    a = np.append(a,r)
            elif isinstance(ai,dict):
                for i in ai:
                    r = recursive_lecture(ai[i])
                    a = np.append(a,r)
            elif isinstance(ai,spm_vec):
                return ai.vec
            return a
        
        def recursive_linking(ai,st,i=0):
            ''' Matches the vector with the pointers in the input '''
            if isinstance(ai,np.ndarray):
                shape = ai.shape
                ai = st[i:i+reduce(operator.mul, shape, 1)]
                ai.shape = shape
                return ai,i+reduce(operator.mul, shape, 1)
            elif isinstance(ai,list):
                l = []
                for j in ai:
                    lt,i = recursive_linking(j,st,i)
                    l += [lt]
                return l,i
            elif isinstance(ai,dict):
                l = {}
                for j in ai:
                    lt,i = recursive_linking(ai[j],st,i)
                    l[j] = lt
                return l,i
            elif isinstance(ai,spm_vec):
                lt,i = recursive_linking(ai.data,st,i)
                return lt,i
            return [],i
        
        self.vec  = np.asarray(recursive_lecture(iterable))
        self.data = recursive_linking(iterable,self.vec)[0]
        
    def __len__(self):
        return len(self.vec)
    
    def __getitem__(self,key):
        # TODO urgent. Make the acces more soffisticated that just accesing
        # the members of the list            
        self.vec[key]            
    def __isub__(self,other):
        if isinstance(other,spm_vec):
            np.subtract(self.vec,other.vec,self.vec)
            return self
        else:
            np.subtract(self.vec,other,self.vec)
            return self
        return self
    def __idiv__(self,other):
        if isinstance(other,spm_vec):
            np.divide(self.vec,other.vec,self.vec)
            return self
        else:
            np.divide(self.vec,other,self.vec)
            return self
    def __imul__(self,other):
        if isinstance(other,spm_vec):
            np.multiply(self.vec,other.vec,self.vec)
            return self
        else:
            np.multiply(self.vec,other,self.vec)
            return self
        return self
    def __iadd__(self,other):
        if isinstance(other,spm_vec):
            np.add(self.vec,other.vec,self.vec)
            return self
        else:
            np.add(self.vec,other,self.vec)
            return self
        return self
    def __add__(self,other):
        g = spm_vec(self.data)
        if isinstance(other,spm_vec):
            np.add(self.vec,other.vec,g.vec)
            return g
        else:
            np.add(self.vec,other,g.vec)
            return g
    
    def __str__(self,values=False):
        s = 'spm_vec:['
        n = 0
        try:
            for i in self.data:
                s += str(type(i)) + ','
                n += 1
                if n > 7:
                    s += '...'
                    break              
        except TypeError:
            pass
        s += ']\n'
        n = 0            
        s += self.vec.__str__()
        return s
                
                  
        
        
        
        

        

class Model:    
    ''' Defines a generative model 
    
    ml = Model length           int
    is_ADEM                     bool
    
    g  = y(t)  = g(x,v,P)     Callable
    gx = dg/dx                Callable
    gv = dg/dv                Callable
    gP = dg/dP                Callable
    f  = dx/dt = f(x,v,P)     Callable
    fx = df/dx                Callable
    fv = df/dv                Callable
    fP = df/dP                Callable
    
    a  = action 
    m  = number of inputs v(i + 1);
    n  = number of states x(i);
    l  = number of output v(i);
    x  = hidden states;
    v  = causal states;
    k  = number of active states
    pE = prior expectation of p model-parameters
    pC = prior covariances of p model-parameters
    hE = prior expectation of h log-precision (cause noise)
    hC = prior covariances of h log-precision (cause noise)
    gE = prior expectation of g log-precision (state noise)
    gC = prior covariances of g log-precision (state noise)
    xC = prior covariances of states
    ph = ?
    pg = ?
    Q  = precision components (input noise)
    R  = precision components (state noise)
    V  = fixed precision (input noise)
    W  = fixed precision (state noise)
    xP = Prior on the states
    sw = Smothness of the states
    sv = Smothness of the causes
    
    E.s   = smoothness (s.d. in time bins)
    E.d   = embedding order q(v)  (i.e., number of derivatives)
    E.n   = embedding order q(x)
    E.dt  = temporal difference
    E.nD = Number of iterations in the D step
    E.nE = Number of iterations in the E step
    E.nM = Number of iterations in the M step
    E.nN = Number of interations in the ???
    
    '''
    def __init__(self,ml):
        ''' ml is the order of the model. This function automatically increases 
        the order of the model to accomodate a ml+1 top level for external 
        causes. All the properties are initiated with lists with empty entries.
        The E field is initiated with a class with empty properties.'''
        
        
        self.ml = ml
        self.is_ADEM  = False
        
        self.a  = [None]*ml
        self.k  = [None]*ml
        self.g  = [None]*ml
        self.gx = [None]*ml
        self.gv = [None]*ml
        self.gP = [None]*ml
        self.f  = [None]*ml
        self.fx = [None]*ml
        self.fv = [None]*ml
        self.fP = [None]*ml
        self.m  = [None]*ml
        self.n  = [None]*ml
        self.l  = [None]*ml
        self.x  = [None]*ml
        self.v  = [None]*ml
        self.pE = [None]*ml
        self.pC = [None]*ml
        self.hE = [None]*ml
        self.hC = [None]*ml
        self.gE = [None]*ml
        self.gC = [None]*ml
        self.xC = [None]*ml
        self.ph = [None]*ml
        self.pg = [None]*ml
        self.Q  = [None]*ml
        self.R  = [None]*ml
        self.V  = [None]*ml
        self.W  = [None]*ml
        self.xP = [None]*ml
        self.sw = [None]*ml
        self.sv = [None]*ml
        
        class E:
            def __init__(selfi):
                selfi.s  = None
                selfi.d  = None
                selfi.n  = None
                selfi.dt = None
                selfi.nD = None
                selfi.nE = None
                selfi.nM = None
                selfi.nN = None
                selfi.linear = None
            def __str__(self,pre='',post='\n'):
                inf = ''
                for i in self.__dict__:
                    inf += '{0}E.{1}:\t{2}{3}'.format(pre,str(i),self.__dict__[i].__str__(),post)
                return inf
                    
                
        self.E = E()
        
    def __str__(self,pre='',post='\n'):
        inf = ''
        for i in self.__dict__:
            if i == 'E':
                inf += self.__dict__[i].__str__(pre=pre,post=post)
            else:
                inf += '{0}{1}:\t{2}{3}'.format(pre,str(i),self.__dict__[i].__str__(),post)
        return inf
                
        
    def add_level(self):
        self.ml += 1
        
        self.a  += [None]
        self.k  += [None]
        self.g  += [None]
        self.gx += [None]
        self.gv += [None]
        self.gP += [None]
        self.f  += [None]
        self.fx += [None]
        self.fv += [None]
        self.fP += [None]
        self.m  += [None]
        self.n  += [None]
        self.l  += [None]
        self.x  += [None]
        self.v  += [None]
        self.pE += [None]
        self.pC += [None]
        self.hE += [None]
        self.hC += [None]
        self.gE += [None]
        self.gC += [None]
        self.xC += [None]
        self.ph += [None]
        self.pg += [None]
        self.Q  += [None]
        self.R  += [None]
        self.V  += [None]
        self.W  += [None]
        self.xP += [None]
        self.sw += [None]
        self.sv += [None]
    
    def check(self):
        ''' Checks fields or reports errors.'''
        
        # order
        #----------------------------------------------------------------------
         
        # set missing fields
        #======================================================================
 
        # check for specification of hidden states
        #----------------------------------------------------------------------
        
        def none_list(l):
            return all([i is None for i in l])
        
        if none_list(self.n) and none_list(self.x): 
            raise ModelError(['n','x'],'Please specify hidden states or their number')
        
        # check supra-ordinate level and add one (with flat priors) if necessary
        #----------------------------------------------------------------------
        
        if hasattr(self.g[self.ml-1], '__call__'):
            self.add_level()
            if not self.m[-2] is None:
                self.l[-1] = self.m[-2]
        self.m[-1] = 0
        self.n[-1] = 0
            
        # default fields for static models (hidden states)
        #----------------------------------------------------------------------
        def r(x=None,v=None,P=None):
                return spr.lil_matrix((1))
                
        if none_list(self.f):
            self.f = [r] * self.ml
            self.x = [np.array(())] * self.ml
            self.n = [0] * self.ml
            
        for i in xrange(self.ml):
            if not hasattr(self.f[i], '__call__'):
                self.f[i] = r
                self.x[i] = np.array(())
                self.n[i] = 0
                
        # consistency and format check on states, parameters and functions
        #======================================================================

        # prior expectation of parameters M.pE
        #----------------------------------------------------------------------
        if none_list(self.pE):
            # Assume fixed parameters
            #------------------------------------------------------------------
            self.pE = [np.array(())] *self.ml
        
        # and priors covariances - p
        #----------------------------------------------------------------------
        if none_list(self.pC):
            # Assume fixed parameters
            #------------------------------------------------------------------
            self.pC = []
            for i in xrange(self.ml):
                p          = np.prod(self.pE[i].shape)
                self.pC    += [np.zeros((p,p))]
                
        # check pC if user specified
        #----------------------------------------------------------------------
        for i in xrange(self.ml):
            # number of parameters
            #------------------------------------------------------------------
            npr  = np.prod(self.pE[i].shape)
 
            # Assume fixed parameters if not specified
            #------------------------------------------------------------------
            
            if self.pC[i] is None:
                self.pC[i] = np.zeros((npr,npr));
 
            # convert variances to covariances if necessary
            #------------------------------------------------------------------
            l = self.pC[i].shape
            
            if max(l) > 1 and max(l)/np.prod(l) == 1: 
                # Sparse diagonal matrix.
                self.pC[i] = np.diag(self.pC[i],0)
            
            # convert variance to covariances if necessary
            #------------------------------------------------------------------
            if np.prod(self.pC[i].shape) == 1 or isinstance(self.pC[i], (int, long, float, complex)):
                # TODO THere is a bug
                self.pC[i] = np.eye(npr) * self.pC[i]
                pass

            # check size
            #------------------------------------------------------------------

            if max(self.pC[i].shape) != npr:
                ModelError('pC','Please check: M.pC[{0}]'.format(i))

        # get inputs
        #----------------------------------------------------------------------
        
        if none_list(self.v):
            v = np.zeros((0,0))
        else:
            v = self.v[-1]
            
        if v is None:
            try:
                v = np.zeros((self.m[-2],1))
            except:
                pass
        
        if v is None:
            try:
                v = np.zeros((self.l[-1],1))
            except:
                pass

        
        self.l[-1]    = len(spm_vec(v).vec)
        self.v[-1]    = v
        
        # check functions
        #----------------------------------------------------------------------
        
        if none_list(self.x):
            self.x = [spr.bsr_matrix((self.n[i],1)) for i in xrange(self.ml)]
        
        for i in [j-1 for j in range(self.ml-1,0,-1)]:
            
            if self.x[i] is None and self.n[i]:
                x = spr.bsr_matrix((self.n[i],1))
            else:
                x = self.x[i]
             
            # check f(x,v,P)
            #-----------------------------------------------------------------
             
            if not hasattr(self.f[i], '__call__'):
                raise ModelError('f','M.f[{0}] is not callable'.format(i))
                 
            try:
                f = self.f[i](x,v,self.pE[i])
            except:
                # TODO bug
                raise ModelError('f','Error while evaluating M.f[{0}].'.format(i))
                pass
             
            if np.prod(x.shape) != np.prod(f.shape):
                raise ModelError('f','Dimensions of x and M.f[{0}](x) are not equal'.format(i))
             

            # check df/dx
            #-----------------------------------------------------------------
                
            if self.fx[i] != None:
               if not hasattr(self.fx[i], '__call__'):
                   raise ModelError('fx','M.fx[{0}] is not callable'.format(i))
               try:
                   self.fx[i](x,v,self.pE[i])
               except:
                   raise ModelError('fx','Error while evaluating M.fx[{0}].'.format(i))

            # check df/dv
            #-----------------------------------------------------------------
                    
            if self.fv[i] != None:
               if not hasattr(self.fv[i], '__call__'):
                   raise ModelError('fv','M.fv[{0}] is not callable'.format(i))
               try:
                   self.fv[i](x,v,self.pE[i])
               except:
                   raise ModelError('fv','Error while evaluating M.fv[{0}].'.format(i))

            # check df/dP
            #-----------------------------------------------------------------
                    
            if self.fP[i] != None:
               if not hasattr(self.fP[i], '__call__'):
                   raise ModelError('fx','M.fx[{0}] is not callable'.format(i))
               try:
                   self.fP[i](x,v,self.pE[i])
               except:
                   raise ModelError('fP','Error while evaluating M.fP[{0}].'.format(i))
                 
            # check g(x,v,P)
            #-----------------------------------------------------------------

            if not hasattr(self.g[i], '__call__'):
                print(i)
                raise ModelError('g','M.g[{0}] is not callable'.format(i))
             
            self.m[i] = np.prod(v.shape)
            try:
                v = self.g[i](x,v,self.pE[i])
            except:
                raise ModelError('g','Error while evaluating M.g[{0}].'.format(i))
            
            self.l[i] = np.prod(v.shape)
            self.n[i] = np.prod(x.shape)
            
            self.v[i] = v
            self.x[i] = x

            # check dg/dx
            #-----------------------------------------------------------------
                
            if self.gx[i] != None:
               if not hasattr(self.gx[i], '__call__'):
                   raise ModelError('gx','M.gx[{0}] is not callable'.format(i))
               try:
                   self.gx[i](x,v,self.pE[i])
               except:
                   raise ModelError('gx','Error while evaluating M.gx[{0}].'.format(i))

            # check dg/dv
            #-----------------------------------------------------------------
                    
            if self.gv[i] != None:
               if not hasattr(self.gv[i], '__call__'):
                   raise ModelError('gv','M.gv[{0}] is not callable'.format(i))
               try:
                   self.gv[i](x,v,self.pE[i])
               except:
                   raise ModelError('gv','Error while evaluating M.gv[{0}].'.format(i))

            # check dg/dP
            #-----------------------------------------------------------------
                    
            if self.gP[i] != None:
               if not hasattr(self.gP[i], '__call__'):
                   raise ModelError('gx','M.gx[{0}] is not callable'.format(i))
               try:
                   self.gP[i](x,v,self.pE[i])
               except:
                   raise ModelError('gP','Error while evaluating M.gP[{0}].'.format(i))
                   
            
            
        # priors on states
        #------------------------------------------------------------------
        
        for i in xrange(self.ml):
            try:
                if np.prod(self.xP[i].shape) == 1:
                    self.xP[i] = np.eye(self.n[i])*self.xP[i]
            except AttributeError:
                    pass
            if self.xP[i] is None:
                self.xP[i] = np.zeros((self.n[i],self.n[i]))
            elif self.xP[i] != (self.n[i],self.n[i]):
                raise ModelError('xP','Please check M.xP[{0}].'.format(i))
        
        # number of x (hidden states)
        #----------------------------------------------------------------------
        
        nx     = sum(self.n)
        
        # check hyperpriors hE - [log]hyper-parameters and components
        #----------------------------------------------------------------------
        
        pP    = 1               # prior precision on log-precisions
        
        for i in xrange(self.ml):    
            # make sure components are cell arrays
            #------------------------------------------------------------------
            
            if self.Q[i] is None:
                self.Q[i] = []
            
            if not isinstance(self.Q[i],list):
                 self.Q[i] = [self.Q[i]]
            
            
            if self.R[i] is None:
                self.R[i] = []
            
            if not isinstance(self.R[i],list):
                 self.R[i] = [self.R[i]]

    
            # check hyperpriors
            #==================================================================
    
            # vectorise
            #------------------------------------------------------------------
            
            self.hE[i] = spm_vec(self.hE[i]).vec
            self.gE[i] = spm_vec(self.gE[i]).vec
    
            # check hyperpriors (expectations)
            #------------------------------------------------------------------
            if self.hE[i] == []:
                self.hE[i] = np.zeros((len(self.Q[i]),1))
                
            if self.gE[i] == []:
                self.gE[i] = np.zeros((len(self.R[i]),1))
    
            # check hyperpriors (covariances)
            #------------------------------------------------------------------
            try:
                # TODO Needs verification
                self.hC[i]*self.hE[i]
                
            except:
                self.hC[i] = np.eye(max(self.hE[i].shape))/pP
            try:
                # TODO Needs verification
                self.gC[i]*self.gE[i]
            except:
                self.gC[i] = np.eye(max(self.gE[i].shape))/pP
            
            if self.hC[i] is None:
                self.hC[i] = np.eye(max(self.hE[i].shape))/pP
            if self.gC[i] is None:
                self.gC[i] = np.eye(max(self.gE[i].shape))/pP
                
            # check Q and R (precision components)
            #==================================================================

    
            # check components and assume i.i.d if not specified
            #------------------------------------------------------------------

            if len(self.Q[i]) > max(self.hE[i].shape):
                self.hE[i] = np.zeros((len(self.Q[i]),1)) + self.hE[i][0]
            
            if len(self.Q[i]) < max(self.hE[i].shape):
                self.Q[i] = [np.eye(self.l[i])]
                self.hE[i] = self.hE[i][0]

            if max(self.hE[i].shape) > max(self.hC[i].shape):
                self.hC[i] = np.eye(max(self.Q[i].shape))*self.hC[i][0]
            
            if len(self.R[i]) > max(self.gE[i].shape):
                self.gE[i] = np.zeros((max(self.R[i].shape),1)) + self.gE[i][0]
                
            if len(self.R[i]) < max(self.gE[i].shape):
                self.R[i] = [np.eye((self.n[i],self.n[i]))]
                self.gE[i] = self.gE[i][0]
            
            if max(self.gE[i].shape) > max(self.gC[i].shape):
                self.gC[i] = np.eye(max(self.R[i].shape))*self.gC[i][0]

            # check consistency and sizes (Q)
            #------------------------------------------------------------------
            for j in xrange(len(self.Q[i])):
                print self.Q[i][j]
                if max(self.Q[i][j].shape) != self.l[i]:
                    raise ModelError('Q','wrong size; M.Q[{0}][{1}]'.format(i,j))
        
            # check consistency and sizes (R)
            #------------------------------------------------------------------
            for j in xrange(len(self.R[i])):
                if max(self.R[i][j].shape) != self.n[i]:
                    raise ModelError('R','wrong size; M.R[{0}][{1}]'.format(i,j))
                    
            
            # check V and W (lower bound on precisions)
            # =================================================================

            # check V and assume unit precision if improperly specified
            #------------------------------------------------------------------
            
            if self.V[i] is None or max(self.V[i].shape) != self.l[i]:
                try:
                    self.V[i] = np.eye(self.l[i])*self.V[i][0]
                except:
                    if self.hE[i] is None and self.ph[i] is None:
                        self.V[i] = np.eye(self.l[i])
                    else:
                        self.V[i] = np.zeros((self.l[i],self.l[i]))
                        
                
            # check W and assume unit precision if improperly specified
            #------------------------------------------------------------------
            if self.W[i] is None or max(self.W[i].shape) != self.n[i]:
                try:
                    self.W[i] = np.eye(self.n[i])*self.W[i][0]
                except:
                    if (self.gE[i] is None or self.gE == []) and self.pg[i] is None:
                        self.W[i] = np.eye(self.n[i])
                    else:
                        self.W[i] = np.zeros((self.n[i],self.n[i]))
        
        # estimation parameters M(1).E.s, n,...
        #======================================================================
        # E.s;                               # smoothness (seconds)
        # E.dt;                              # time step
        # E.d;                               # approximation order of q(x,v)
        # E.n;                               # order of embedding (n >= d)
 
         # temporal smoothness - s.d. of kernel
         #---------------------------------------------------------------------
         
        if not self.E.s is None:
            if nx:
                self.E.s = 0.5
            else:
                self.E.s = 0
        # time step
        #----------------------------------------------------------------------
        if not self.E.dt is None:
            self.E.dt = 1
        
        # embedding orders
        #----------------------------------------------------------------------
        if not self.E.dt is None:
            if nx:
                self.E.d = 2
            else:
                self.E.d = 0
        if not self.E.n is None:
            if nx:
                self.E.n = 6
            else:
                self.E.n = 0
        
        self.E.d = min(self.E.d,self.E.n)

        # number of iterations
        #----------------------------------------------------------------------
        if not self.E.nD is None:
            if nx:
                self.E.nD = 1
            else: self.E.nD = 8

        if not self.E.nE is None:
            self.E.nE = 8

        if not self.E.nM is None:
            self.E.nM = 8

        if not self.E.nN is None:
            self.E.nN = 8
            
            
        # checks on smoothness hyperparameter
        #----------------------------------------------------------------------
        
        for i in xrange(self.ml):
            if self.sv[i] is None:
                self.sv[i] = self.E.s
            if self.sw[i] is None:
                self.sw[i] = self.E.s
        
        # check on linear approximation scheme
        #----------------------------------------------------------------------
        
        if self.E.linear is None:
            self.E.linear = False

        
        # checks on estimability
        #----------------------------------------------------------------------
 
        # check that there are informative priors on the states or the causes
        #----------------------------------------------------------------------
        #Q     = ~norm(M(end).V,1);
        #for i = 1:(g - 1)
        #    P = norm(M(i).pC,1) > exp(8);
        #    if P && Q
        #        warndlg('please use informative priors on causes or parameters')
        #    end
        #end


class spm_ADEM:
    
    def __init__(self,M,G,C,U=None,mode=True):
        ''' Initilizes a model '''

        M.check()
        self.M = M

        # check data or generative model
        # ---------------------------------------------------------------------
        
        N         = max(C.shape)
        self.mode = mode
        
        # Default priors
        # ------------------------------------------------------------------------- 
        
        if U is None:
            self.U = np.zeros((M.l[-1],N))
        
        
        
        
        
        
        
        
    
    pass


            
if __name__ == '__init__':

    M = Model(1)
    
    def f(x,v,P):
        return np.array([0])
    def g(x,v,P):
        return np.array([0])
        
    M.f[0] = f
    M.g[0] = g
    M.x[0] = np.array([0])
    
    M.check()
    
    g = np.array((1,2,34))
    f = np.array((45,6,7))
    l = [g,f]
    n = {'p':l,'r':np.array([1,2])}
    e = spm_vec(n)
    

                   
            


# To implement:
    
    # spm_DEM_embed
    # spm_speye
    # spm_ADEM_diff
    # kron
    # Implement a data structure that solves the problem with spm_vec and spm_unvec
    # spm_DEM_eval
    # spm_logdet
    # spm_cat
    # spm_dx


# Implements the D step of Active Dynamical Expectation Maximazation.

